import pymongo
import matplotlib.pyplot as plt

client = pymongo.MongoClient(host='127.0.0.1', port=27017)
# db 表示数据库名称，username 用户名 password 密码
#client.db.authenticate(useranme, password)
# 建立数据库连接
db = client.CNNVD_vul # or db = client['test']
collection = db.vuldb # or collection = db['students']

#各月份漏洞数量统计图
def month_display(year):
    #monthname记录各月份名称，monthlist记录各月份的漏洞数量
    monthname=[]
    monthlist=[]
    #循环得到各月份的漏洞数量
    for i in range(1,13):
        #月份记录
        time=str(i)+'月'
        monthname.append(time)
        #起始时间time_start和终止时间time_end
        time_s=year+"-"+"%02d"%(i)+"-00"
        time_e=year+"-"+"%02d"%(i)+"-40"
        #匹配搜索该月份内的漏洞
        res = collection.find({"publishTime":
                        {
                                '$gte':time_s,
                                "$lte":time_e
                        }
                        })
        lres=list(res)
        #统计该月份的漏洞数量
        #print(len(lres))
        monthlist.append(len(lres))

    #plt函数套公式使用
    plt.subplot(2, 2, 1)
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
    plt.rcParams['axes.unicode_minus'] = False
    plt.title(year+"年各月份漏洞数量统计图")
    plt.bar(monthname, monthlist)

#各危害等级漏洞数量统计图
def hazardlevel_display(year):
    #漏洞危害等级表，以及其对应的漏洞数量表
    level=["暂未定级","超危","高危","中危","低危"]
    vullist=[]
    #起始时间time_start和终止时间time_end
    time_s=year+"-01"
    time_e=year+"-12"
    #按照漏洞危害等级循环，得到相应危害等级的漏洞的数量
    for i in range(0,5):
        res = collection.find({'hazardLevel': i,"publishTime":{'$gte':time_s,"$lte":time_e}})
        lres=list(res)
        #统计数量
        #print(len(lres))
        vullist.append(len(lres))

    #plt函数套公式使用
    plt.subplot(2, 2, 2)
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
    plt.rcParams['axes.unicode_minus'] = False
    plt.title(year+"年各危险等级漏洞")
    plt.bar(level, vullist)

#各类型漏洞数量统计图
def typeName_display(year):
    #起始时间time_start和终止时间time_end   
    time_s=year+"-01"
    time_e=year+"-12"
    #得到数据库中所有的typename, 格式： {'typeName': '信息泄露'}
    res = collection.find({"publishTime":{'$gte':time_s,"$lte":time_e}},{"varchar1":1,'_id':0})
    lres=list(res)
    #列表中的元素都是如上格式的字典，我们将其转化为字符串，然后使用set（集合）去重
    for i in range(len(lres)):
        lres[i]=str(lres[i])
    #s1中包含了该时间段内的所有typeName类型，且不重复
    s1=set(lres)
    #记录数据库中每个typeName的名字ltypename，及其包含的漏洞数lnumber
    ltypename=[]
    lnumber=[]
    #后续进行数据库查询的筛选条件
    dtime={"publishTime":{'$gte':time_s,"$lte":time_e}}
    for i in s1:
        #将之前转化为字符串的{'typeName': '信息泄露'}，再转化为字典
        i=eval(i)
        #加入该typeName的名字，如果为{'typeName': '信息泄露'}，则加入字符串'信息泄露'
        ltypename.append(str(i['varchar1']))
        #确定i为字典类型，才可以调用字典的函数
        i=dict(i)
        #扩充字典i，使其格式为{'typeName': '信息泄露',"publishTime":{'$gte':time_s,"$lte":time_e}}
        #这样可以直接通过查找功能，找到数据库中该时间段该typeName有多少漏洞
        i.update(dtime)
        #查找
        res = collection.find(i)
        lres=list(res)
        #print(i['typeName'],len(lres))
        #匹配到很多漏洞，而我们这里只需要一个总数，直接append(len(lres))。
        lnumber.append(len(lres))
    #冒泡排序,依据lnumber，同时对ltypename和lnumber排序
    combined = list(zip(lnumber, ltypename))
    # 使用sorted函数对combined进行排序，排序依据是元组的第一个元素（即第一个列表的值）
    sorted_combined = sorted(combined, key=lambda x: x[0],reverse=True)
    # 解压排序后的列表
    lnumber, ltypename = zip(*sorted_combined)
    #取排序后数量最大的前20位
    lnumber=lnumber[:20:1]
    ltypename=ltypename[:20:1]
    
    #plt函数套公式使用
    plt.subplot(2, 2, 3)
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
    plt.rcParams['axes.unicode_minus'] = False
    plt.title(year+"年各类型漏洞TOP20")
    plt.barh(ltypename, lnumber)


#各affectedVendor(受影响的供应商)漏洞数量统计图
def affectedVendo_display(year):
#得到所有类型的affectedVendor(受影响的供应商)：    
    #起始时间time_start和终止时间time_end，后面的表示从1月到12月
    time_s=year+"-01"
    time_e=year+"-12"
    #得到数据库中所有的affectedVendor, 格式：{'affectedVendor': '个人开发者'}
    res = collection.find({"publishTime":{'$gte':time_s,"$lte":time_e}},{"affectedVendor":1,'_id':0})
    lres=list(res) 
    #列表中的元素都是如上格式的字典，我们将其转化为字符串，然后使用set（集合）去重
    for i in range(len(lres)):
        lres[i]=str(lres[i])
    #s1中包含了该时间段内的所有affectedVendor类型，且不重复
    s1=set(lres)

#得到所有类型的affectedVendor,及其漏洞数量：
    #记录数据库中每个affectedVendor的名字lvname，及其包含的漏洞数lvnumber
    lvname=[]
    lvnumber=[]
    #后续进行数据库查询的筛选条件
    dtime={"publishTime":{'$gte':time_s,"$lte":time_e}}
    #遍历所有的affectedVendor，并得到其包含的漏洞
    for i in s1:
        #将之前转化为字符串的{'affectedVendor': '个人开发者'}，再转化为字典
        i=eval(i)
        #加入该affectedVendor的名字，如果为{'affectedVendor': '个人开发者'}，则加入字符串'个人开发者'
        lvname.append(str(i['affectedVendor']))
        #确定i为字典类型，才可以调用字典的函数
        i=dict(i)
        #扩充字典i，使其格式为{'affectedVendor': '个人开发者',"publishTime":{'$gte':time_s,"$lte":time_e}}
        #这样可以直接通过查找功能，找到数据库中该时间段该affectedVendor有多少漏洞
        i.update(dtime)
        #查找
        res = collection.find(i)
        lres=list(res)
        #print(i['affectedVendor'],len(lres))
        #匹配到很多漏洞，而我们这里只需要一个总数，直接append(len(lres))。
        lvnumber.append(len(lres))
    
    #冒泡排序,依据lnumber，同时对ltypename和lnumber排序
    combined = list(zip(lvnumber, lvname))
    # 使用sorted函数对combined进行排序，排序依据是元组的第一个元素（即第一个列表的值）
    sorted_combined = sorted(combined, key=lambda x: x[0],reverse=True)
    # 解压排序后的列表
    lvnumber, lvname = zip(*sorted_combined)
    #取排序后数量最大的前20位
    lvnumber=lvnumber[:20:1]
    lvname=lvname[:20:1]

    #plt函数套公式使用
    plt.subplot(2, 2, 4)
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
    plt.rcParams['axes.unicode_minus'] = False
    plt.title(year+"年各个公司(组织)漏洞top20")
    plt.barh(lvname, lvnumber)



if __name__ == '__main__':
    # 查找重复项
    pipeline = [
        {"$group": {"_id": "$cnnvdCode", "duplicates": {"$push": "$_id"}, "count": {"$sum": 1}}},
        {"$match": {"count": {"$gt": 1}}}
    ]
    duplicate_groups = list(collection.aggregate(pipeline))

    # 删除重复项
    for group in duplicate_groups:
        # 获取重复项的ID列表
        duplicate_ids = group["duplicates"][1:]
        # 删除重复项
        collection.delete_many({"_id": {"$in": duplicate_ids}})
        print(f"Deleted {len(duplicate_ids)} duplicates.")
    month_display("2022")
    print("图一完成！")
    hazardlevel_display("2023")
    print("图二完成！")
    typeName_display("2022")
    print("图三完成！")
    affectedVendo_display("2022")
    print("图四完成！")
    #整合展示
    plt.show()
    # 关闭数据库连接
    client.close()

